
import tensorflow as tf
import os
from tensorflow.keras.models import model_from_json
from tensorflow.keras.models import load_model
# from keras.backend import identity
# from keras.backend import concatenate
from tensorflow import identity
from tensorflow import concat
import json
from faster_rcnn.layers.anchors import Anchor
from faster_rcnn.layers.batch_norm import BatchNorm
from faster_rcnn.layers.proposals import RpnToProposal
from faster_rcnn.layers.target import RpnTarget
from faster_rcnn.layers.target import DetectTarget
from faster_rcnn.layers.roi_align import RoiAlign
from faster_rcnn.layers.detect_boxes import ProposalToDetectBox
from faster_rcnn.layers.clip_boxes import ClipBoxes
from faster_rcnn.layers.losses import rpn_cls_loss, rpn_regress_loss, detect_regress_loss, detect_cls_loss
from faster_rcnn.layers.specific_to_agnostic import deal_delta


custom_objects = {
    "Anchor": Anchor, "RpnToProposal": RpnToProposal, "RpnTarget": RpnTarget,
    "DetectTarget": DetectTarget, "RoiAlign": RoiAlign,
    "ProposalToDetectBox": ProposalToDetectBox, "ClipBoxes": ClipBoxes,
    "rpn_cls_loss": rpn_cls_loss, "rpn_regress_loss": rpn_regress_loss,
    "detect_regress_loss": detect_regress_loss, "detect_cls_loss": detect_cls_loss,
    "deal_delta": deal_delta, "identity": identity, "concat": concat}


def json_load_model(model_weight_path, model_json_path):
    with open(model_json_path, 'rb') as f:
        model_json = json.load(f)

    model = model_from_json(model_json, custom_objects=custom_objects)
    model.load_weights(model_weight_path, by_name=True)

    print("weight len: {}".format(len(model.get_weights())))
    # print("weight[1]: {}".format(model.get_weights()[1]))
    # print("weight[323]: {}".format(model.get_weights()[323]))
    return model


def all_load_model(model_weight_path):
    model = load_model(model_weight_path, custom_objects=custom_objects)

    print("input.size: {}".format(len(model.inputs)))
    print("output.size: {}".format(len(model.outputs)))
    return model


# save keras models to pb
def tf_save_model(load_model_type, model_paths, pb_save_path):
    if load_model_type == 1:
        model = all_load_model(model_paths[0])
    else:
        model = json_load_model(model_paths[0], model_paths[1])
    model_input1 = tf.saved_model.utils.build_tensor_info(model.inputs[0])
    model_input2 = tf.saved_model.utils.build_tensor_info(model.inputs[1])

    model_output1 = tf.saved_model.utils.build_tensor_info(model.outputs[0])
    model_output2 = tf.saved_model.utils.build_tensor_info(model.outputs[1])
    model_output3 = tf.saved_model.utils.build_tensor_info(model.outputs[2])
    model_output4 = tf.saved_model.utils.build_tensor_info(model.outputs[3])
    model_output5 = tf.saved_model.utils.build_tensor_info(model.outputs[4])

    prediction_signature = (
        tf.saved_model.signature_def_utils.build_signature_def(
            inputs={"input_image": model_input1,
                    "input_image_meta": model_input2},
            outputs={"detect_boxes": model_output1, "class_scores": model_output2,
                     "detect_class_ids": model_output3, "detect_class_logits": model_output4,
                     "image_meta": model_output5},
            method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME))

    builder = tf.saved_model.builder.SavedModelBuilder(pb_save_path)
    with tf.keras.backend.get_session() as sess:
        sess.run(tf.global_variables_initializer())
        builder.add_meta_graph_and_variables(
            sess=sess, tags=[tf.saved_model.tag_constants.SERVING],
            signature_def_map={
                "serving_default": prediction_signature})
        builder.save()


def tf_contrib_save_model(load_model_type, model_paths, pb_save_path):
    if load_model_type == 1:
        model = all_load_model(model_paths[0])
    else:
        model = json_load_model(model_paths[0], model_paths[1])
    tf.contrib.saved_model.save_keras_model(model, pb_save_path)


model_paths = ["../models/{}all_info.h5".format("predict_"), '../models/{}models.json'.format("predict_")]
print("exist model_weight_path: {}".format(os.path.exists(model_paths[0])))


save_model_type = 2
load_model_type = 2
if save_model_type == 1:
    model_save_path = "../models/{}pb".format("predict_tf_")
    tf_save_model(load_model_type, model_paths, model_save_path)
else:
    model_save_path = "../models/{}pb".format("predict_contrib_")
    tf_contrib_save_model(load_model_type, model_paths, model_save_path)

